import functools
from pyson.types import is_object, is_array, is_container
from pyson.iface.iterators import iteritems
from copy import copy as copy_, deepcopy as deepcopy_

__all__ = ['copy', 'deepcopy', 'empty_from', 'register_constructor', 'get_constructor']

def copy(obj, preserve_type=False):
    '''Return a dict or a list with a shallow copy of object. If 
    ``preserve_type`` is True, it tries to return a copy using the same type 
    as the input object.
    
    Examples
    --------
    
    The default behavior is to return lists or dicts
    
    >>> copy((1, 2, 3))
    [1, 2, 3]
    
    If the argument ``preserve_type`` is True, it tries to return a copy with 
    the same type as ``obj``.
    
    >>> copy((1,2,3), preserve_type=True)
    (1, 2, 3)
    
    '''

    if preserve_type:
        return copy_preserve(obj)
    else:
        if is_object(obj):
            return dict(obj)
        elif is_array(obj):
            return list(obj)

    raise TypeError('object of type %s is not a container' % type(obj))

def empty_from(obj, fill_in=None, preserve_type=False):
    '''Return a list or dictionary with the same keys of obj and filled with 
    ``fill_in`` values.
    
    Examples
    --------
    
    >>> empty_from((1,2,3))
    [None, None, None]
    
    >>> empty_from({'foo': 'bar', 'ham': 'spam'}, 0)
    {'foo': 0, 'ham': 0}
    '''

    if is_object(obj):
        data = dict((k, fill_in) for k in obj)
    elif is_array(obj):
        data = [fill_in] * len(obj)
    else:
        raise TypeError('object of type %s is not a container' % type(obj))

    if preserve_type:
        return get_constructor(type(obj))(data)
    else:
        return data

CONSTRUCTORS = {}
def register_constructor(tt, function=None, overwrite=False):
    '''
    Register a constructor ``function`` for the given type ``tt``.
    
    Constructors functions for *object* types take a dictionary 
    (or a ``tt`` instance) as the single positional argument. Likewise, *array*
    constructors take a single list (or tt) instance as argument.
    
    Parameters
    ----------
    tt : type
        Type to be registered.
    function : callable
        Function that performs the conversion. 
    overwrite : bool
        Forces overwrite of a previous constructor for the type ``tt`` if it 
        exists. 
    '''
    # Decorator version of register call
    if function is None:
        return functools.partial(register_constructor, tt)

    if not isinstance(tt, type):
        raise TypeError("'tt' must be a type, got %s" % type(tt))
    if tt in CONSTRUCTORS and not overwrite:
        raise ValueError("type %s is already registered" % tt)

    CONSTRUCTORS[tt] = function

def get_constructor(obj, lazy=True):
    '''Return the constructor function for object.'''

    if lazy:
        return CONSTRUCTORS.get(type(obj), copy_preserve)
    else:
        try:
            return CONSTRUCTORS[type(obj)]
        except KeyError:
            raise TypeError('there is no constructor for type %s' % type(obj))

def copy_preserve(obj):
    '''Return a copy of obj preserving its type.'''

    try:
        tt = CONSTRUCTORS[type(obj)]
    except KeyError:
        if is_object(obj):
            return obj.copy()
        elif is_array(obj):
            return obj[:]
        else:
            raise TypeError('object of type %s is not a container' % type(obj))
    else:
        return tt(obj)

for tt in [list, dict, tuple, set, frozenset]:
    register_constructor(tt, tt)

COPY_FUNCTIONS = {'copy': copy_, 'deepcopy': deepcopy_}
def deepcopy(obj, preserve_type=False, copy='copy'):
    if isinstance(copy, basestring):
        copy = COPY_FUNCTIONS.get(copy, copy)
    if is_container(obj):
        new = empty_from(obj)
        for k, v in iteritems(obj):
            new[k] = deepcopy(v, preserve_type, copy)
        return new
    else:
        return copy(obj)

if __name__ == '__main__':
    import doctest
    doctest.testmod(optionflags=doctest.REPORT_ONLY_FIRST_FAILURE, verbose=0)
